# coding=utf-8
import os
import cv2
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import time
import glob


att_columns = [u'neckline_left', u'neckline_right',
               u'center_front', u'shoulder_left', u'shoulder_right', u'armpit_left',
               u'armpit_right', u'waistline_left', u'waistline_right', u'cuff_left_in',
               u'cuff_left_out', u'cuff_right_in', u'cuff_right_out', u'top_hem_left',
               u'top_hem_right', u'waistband_left', u'waistband_right',
               u'hemline_left', u'hemline_right', u'crotch', u'bottom_left_in',
               u'bottom_left_out', u'bottom_right_in', u'bottom_right_out']

att_list = {
    "blouse":{"droplist":[u'waistline_left', u'waistline_right',u'waistband_left', u'waistband_right',
                          u'hemline_left', u'hemline_right', u'crotch', u'bottom_left_in',
                          u'bottom_left_out', u'bottom_right_in', u'bottom_right_out'],
              "joints_list" : [u'neckline_left', u'neckline_right',
                               u'center_front', u'shoulder_left', u'shoulder_right', u'armpit_left',
                               u'armpit_right', u'cuff_left_in', u'cuff_left_out', u'cuff_right_in',
                               u'cuff_right_out', u'top_hem_left', u'top_hem_right'],
              "num_keypoints":13},
    "dress":{"droplist":[u'top_hem_left',u'top_hem_right', u'waistband_left', u'waistband_right', u'crotch', u'bottom_left_in',
                         u'bottom_left_out', u'bottom_right_in', u'bottom_right_out'],
             "joints_list": [u'neckline_left', u'neckline_right', u'center_front', u'shoulder_left', u'shoulder_right', u'armpit_left',
                             u'armpit_right', u'waistline_left', u'waistline_right', u'cuff_left_in', u'cuff_left_out', u'cuff_right_in', u'cuff_right_out', u'hemline_left', u'hemline_right'],
             "num_keypoints": 15},
    "outwear":{"droplist":[u'center_front',  u'waistband_left', u'waistband_right',u'hemline_left', u'hemline_right', u'crotch', u'bottom_left_in',
                           u'bottom_left_out', u'bottom_right_in', u'bottom_right_out'],
               "joints_list": [u'neckline_left', u'neckline_right', u'shoulder_left', u'shoulder_right', u'armpit_left', u'armpit_right',
                               u'waistline_left', u'waistline_right', u'cuff_left_in', u'cuff_left_out', u'cuff_right_in', u'cuff_right_out', u'top_hem_left', u'top_hem_right'],
               "num_keypoints": 14},
    "skirt":{"droplist":[u'neckline_left', u'neckline_right', u'shoulder_left', u'shoulder_right', u'armpit_left',
                         u'armpit_right', u'waistline_left', u'waistline_right', u'cuff_left_in',u'cuff_left_out', u'cuff_right_in',
                         u'cuff_right_out', u'top_hem_left',u'top_hem_right',u'crotch', u'bottom_left_in',u'bottom_left_out',
                         u'bottom_right_in', u'bottom_right_out',u'center_front'],
             "joints_list":[u'waistband_left', u'waistband_right',u'hemline_left', u'hemline_right'],
             "num_keypoints": 4},
    "trousers":{"droplist":[u'neckline_left', u'neckline_right', u'center_front', u'shoulder_left', u'shoulder_right', u'armpit_left',
                            u'armpit_right', u'waistline_left', u'waistline_right', u'cuff_left_in', u'cuff_left_out', u'cuff_right_in', u'cuff_right_out', u'top_hem_left',
                            u'top_hem_right',u'hemline_left', u'hemline_right'],
                "joints_list": [u'waistband_left', u'waistband_right', u'crotch', u'bottom_left_in', u'bottom_left_out', u'bottom_right_in', u'bottom_right_out'],
                "num_keypoints": 7}
}


def center_one_image(path_img, category, coordination, save_root, extend_x=0.1, extend_y=0.1, debug=False):
    """
    单张图片中心化
    :param path_img:
    :param coordition:
    :return:
    """
    filename = path_img.split('/')[-1]
    img = cv2.imread(path_img)

    joints_list = att_list[category]['joints_list']
    valid_coord = []
    for joint_name in joints_list:
        index = att_columns.index(joint_name)
        valid_coord.append(coordination[index])

    # 去掉-1,防止数据集标注错误
    ind_valid = []
    for idx in range(len(valid_coord)):
        x_y_v = valid_coord[idx]
        if x_y_v[2] == -1:
            continue
        ind_valid.append(idx)
    valid_coord = [valid_coord[idx] for idx in ind_valid]

    min_x = np.min([valid_coord[idx][0] for idx in range(len(valid_coord))])
    min_y = np.min([valid_coord[idx][1] for idx in range(len(valid_coord))])
    max_x = np.max([valid_coord[idx][0] for idx in range(len(valid_coord))])
    max_y = np.max([valid_coord[idx][1] for idx in range(len(valid_coord))])
    paded_width = 25   # int((max_x - min_x) * extend_x * 0.5)
    paded_height = 25  # int((max_y - min_y) * extend_y * 0.5)
    x = max(min_x - paded_width, 0)
    y = max(min_y - paded_height, 0)
    width = max_x - min_x + paded_width * 2
    height = max_y - min_y + paded_height * 2
    if (x + width) > img.shape[1]:
        width = img.shape[1] - x - 1
    if (y + height) > img.shape[0]:
        height = img.shape[0] - y - 1

    # # 保持原始长宽比
    # rate_ori = img.shape[0] * 1.0 / img.shape[1]
    # rate_now = height * 1.0 / width

    bbox = [x, y, width, height]

    if not os.path.exists(os.path.join(save_root, 'Images')):
        os.mkdir(os.path.join(save_root, 'Images'))
    if not os.path.exists(os.path.join(save_root, 'Images', category)):
        os.mkdir(os.path.join(save_root, 'Images', category))
    save_path_img = os.path.join(save_root, 'Images', category, filename)
    cv2.imwrite(save_path_img, img[y: y + height, x: x + width, :])
    # time.sleep(0.2)

    if debug:
        for idx in range(len(valid_coord)):
            x_y_v = valid_coord[idx]
            cv2.circle(img, (x_y_v[0], x_y_v[1]), 3, (0, 0, 255), 2)
        cv2.rectangle(img, (x, y), (x + width, y + height), (0, 0, 255), 2)
        plt.figure()
        plt.imshow(img)
        plt.show()

    return bbox, img, save_path_img


def check_save_images(path_root):
    images_path_list = glob.glob(os.path.join(path_root, '*.jpg'))
    for idx in range(len(images_path_list)):
        img = cv2.imread(images_path_list[idx])
        try:
            shape = img.shape
        except:
            print images_path_list[idx]

        if len(shape) != 3:
            print(images_path_list[idx])


if __name__ == '__main__':
    path_root = '/media/hszc/data1/FaishionAI/round2'
    save_root = '/media/hszc/data1/FaishionAI/round2_crop/val'
    save_csv_name = 'GCN512_bs16_3.98_attr_val.csv'

    if not os.path.exists(save_root):
        os.makedirs(save_root)

    # csv_path = os.path.join(path_root, 'Annotations', save_csv_name)
    csv_path = '/home/hszc/zhangchi/Landmarks/attr_to_crop/GCN512_bs16_3.98_attr_val.csv'
    csv_content = pd.read_csv(csv_path)
    image_files_list = csv_content['image_id']
    image_categories = csv_content['image_category']
    new_image_files_list = []
    print csv_content
    all_coords = []
    all_bboxes = []
    for idx in range(len(image_files_list)):
        path_img = os.path.join(path_root, image_files_list[idx])
        print(path_img)
        coord = []
        for att in att_columns:
            x, y, v = csv_content.iloc[idx][att].split('_')
            coord.append([int(x), int(y), int(v)])
        # bbox = center_one_image(path_img, image_categories[idx], coord, save_root, 0.12, 0.12, False)
        # 修改 2018.05.08 扩大bbox
        bbox, img,new_path = center_one_image(path_img, image_categories[idx], coord, save_root, debug=False)
        new_image_files_list.append(new_path)
        all_bboxes.append(bbox)
        # 新的图片对应的新的坐标
        new_cood = []
        droplist = att_list[image_categories[idx]]['droplist']
        for new_j in range(len(coord)):
            if att_columns[new_j] in droplist:
                new_cood.append([-1, -1, -1])
                continue
            if coord[new_j][2] == -1:
                new_cood.append(coord[new_j])
            else:
                old_xyv = coord[new_j]

                # 修改 2018.05.08 检查点是否被crop掉了
                x_new = old_xyv[0] - bbox[0]
                y_new = old_xyv[1] - bbox[1]
                if x_new < 0 or y_new < 0 or x_new > bbox[2] or y_new > bbox[3]:
                    print('×*×*×*×*×*×*×*×*×*×*×*×*×**×*×*×*shit, key point is cropped: {}'.format(att_columns[new_j]))
                    new_cood.append([-1, -1, -1])
                    # cv2.circle(img, (old_xyv[0], old_xyv[1]), 5, (255, 0, 0), 2)
                    # plt.figure()
                    # plt.imshow(img)
                    # plt.show()
                else:
                    new_cood.append([x_new, y_new, old_xyv[2]])
        new_cood = [str(each_coord[0]) + '_' + str(each_coord[1]) + '_' + str(each_coord[2]) for each_coord
                    in new_cood]
        all_coords.append(new_cood)

    # 保存为新的annotations.csv
    if not os.path.exists(os.path.join(save_root, 'Annotations')):
        os.mkdir(os.path.join(save_root, 'Annotations'))
    save_data = {}
    save_data['image_id'] = new_image_files_list
    save_data['image_category'] = image_categories
    for idx in range(len(att_columns)):
        save_data[att_columns[idx]] = [one_row[idx] for one_row in all_coords]
    df_data = pd.DataFrame(save_data, columns=['image_id', 'image_category'] + att_columns)
    df_data.to_csv(os.path.join(save_root, 'Annotations', save_csv_name), index=False)

    # # bbox 保存,最终结果做准备
    # save_bbox = {}
    # col_ind = ['x', 'y', 'width', 'height']
    # save_bbox['image_id'] = image_files_list
    # save_bbox['image_category'] = image_categories
    # for idx in range(4):
    #     save_bbox[col_ind[idx]] = [one_row[idx] for one_row in all_bboxes]
    # df_bbox = pd.DataFrame(save_bbox, columns=['image_id', 'image_category'] + col_ind)
    # df_bbox.to_csv(os.path.join(save_root, 'bbox.csv'), index=False)
    #
    # # class_yifu = ['blouse', 'dress', 'outwear', 'skirt', 'trousers']
    # # for each in class_yifu:
    # #     check_save_images('/media/yzl/data/dataset/fashion/round2/Images/' + each)
    #
